import torch
import torchvision
import torch.nn as nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter("log3")
dataset = torchvision.datasets.CIFAR10("./datavision",train=False,transform=torchvision.transforms.ToTensor(),
                                       download=True)
dataloader = DataLoader(dataset,batch_size=64,shuffle=True)

class Module(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(in_channels=3,out_channels=6,kernel_size=(3,3),stride=1,padding=0)

    def forward(self,x):
        x = self.conv(x)
        return x


module = Module()
for epoch in range(3):
    step = 0
    for data in dataloader:
        imgs,targets = data
        output = torch.reshape(module(imgs),[-1,3,30,30])
        writer.add_images(f"{epoch}",output,step)
        step += 1
writer.close()
